import paddle
import paddleseg.transforms as T
import numpy as np
from paddleseg.core import infer

def nn_infer(model, img_path, model_path):
    # 网络定义
    para_state_dict = paddle.load(model_path)
    model.set_dict(para_state_dict)
    # 预测结果
    transforms = T.Compose([
        T.Resize(target_size=(256, 256)),
        T.Normalize()
    ])
    img, _ = transforms(img_path)
    img = paddle.to_tensor(img[np.newaxis, :])
    pre = infer.inference(model, img)
    pred = paddle.argmax(pre, axis=1).numpy().reshape((256, 256))
    return pred.astype('uint8')